import colorsys
import os
import time

import numpy as np
import torch
import torch.nn as nn
from PIL import Image, ImageDraw, ImageFont
from thop import clever_format, profile
from torchsummary import summary

from utils.data import cvtColor, get_new_img_size, resize_image, preprocess_input
from settings import set_detector
from utils.model import load_model_weights

class DETECTOR(object):
    def __init__(self, args):

        self.device = args.device
        self.cuda   = args.cuda

        self.input_shape = args.input_shape
        self.class_names = args.class_names
        self.num_classes = args.num_classes

        args.detector, _, _, args.decoder  = set_detector(args)

        self.bbox_util  = args.decoder
        self.nms_iou = args.nms_iou
        self.confidence = args.confidence
            
        load_model_weights(args)
        
        self.model = args.detector.to(args.device)
        if len(args.gpu_id)>1:
            self.model = nn.DataParallel(self.model)

        #---------------------------------------------------#
        #   画框设置不同的颜色
        #---------------------------------------------------#
        hsv_tuples = [(x / self.num_classes, 1., 1.) for x in range(self.num_classes)]
        self.colors = list(map(lambda x: colorsys.hsv_to_rgb(*x), hsv_tuples))
        self.colors = list(map(lambda x: (int(x[0] * 255), int(x[1] * 255), int(x[2] * 255)), self.colors))
    
    #---------------------------------------------------#
    #   检测图片
    #---------------------------------------------------#
    def detect_image(self, image, crop = False, count = False):
        #---------------------------------------------------#
        #   计算输入图片的高和宽
        #---------------------------------------------------#
        image_shape = np.array(np.shape(image)[0:2])
        #---------------------------------------------------#
        #   计算resize后的图片的大小，resize后的图片短边为600
        #---------------------------------------------------#
        # input_shape = get_new_img_size(image_shape[0], image_shape[1])
        #---------------------------------------------------------#
        #   在这里将图像转换成RGB图像，防止灰度图在预测时报错。
        #   代码仅仅支持RGB图像的预测，所有其它类型的图像都会转化成RGB
        #---------------------------------------------------------#
        image       = cvtColor(image)
        #---------------------------------------------------------#
        #   给原图像进行resize，resize到短边为600的大小上
        #---------------------------------------------------------#
        image_data  = resize_image(image, [self.input_shape[1], self.input_shape[0]])
        #---------------------------------------------------------#
        #   添加上batch_size维度
        #---------------------------------------------------------#
        image_data  = np.expand_dims(np.transpose(preprocess_input(np.array(image_data, dtype='float32')), (2, 0, 1)), 0)

        with torch.no_grad():
            images = torch.from_numpy(image_data)
            if self.cuda:
                images = images.to(self.device)
            
            #-------------------------------------------------------------#
            #   roi_cls_locs  建议框的调整参数
            #   roi_scores    建议框的种类得分
            #   rois          建议框的坐标
            #-------------------------------------------------------------#
            #TODO
            # roi_cls_locs, roi_scores, rois, _ = self.model(images=images, mode='inference')
            output = self.model(images)
            #-------------------------------------------------------------#
            #   利用classifier的预测结果对建议框进行解码，获得预测框
            #-------------------------------------------------------------#
            # results = self.bbox_util.forward(roi_cls_locs, roi_scores, rois, image_shape, input_shape, 
            #                             nms_iou = self.nms_iou, confidence = self.confidence)
            results = self.bbox_util.forward(output, image_shape, self.input_shape, 
                                                    nms_iou = self.nms_iou, confidence = self.confidence)
            #---------------------------------------------------------#
            #   如果没有检测出物体，返回原图
            #---------------------------------------------------------#           
            if len(results[0]) <= 0:
                return image
                
            top_label   = np.array(results[0][:, 5], dtype = 'int32')
            top_conf    = results[0][:, 4]
            top_boxes   = results[0][:, :4]
        
        #---------------------------------------------------------#
        #   设置字体与边框厚度
        #---------------------------------------------------------#
        font        = ImageFont.truetype(font='model_data/simhei.ttf', size=np.floor(3e-2 * image.size[1] + 0.5).astype('int32'))
        thickness   = int(max((image.size[0] + image.size[1]) // np.mean(self.input_shape), 1))
        #---------------------------------------------------------#
        #   计数
        #---------------------------------------------------------#
        if count:
            print("top_label:", top_label)
            classes_nums    = np.zeros([self.num_classes])
            for i in range(self.num_classes):
                num = np.sum(top_label == i)
                if num > 0:
                    print(self.class_names[i], " : ", num)
                classes_nums[i] = num
            print("classes_nums:", classes_nums)
        #---------------------------------------------------------#
        #   是否进行目标的裁剪
        #---------------------------------------------------------#
        if crop:
            for i, c in list(enumerate(top_label)):
                top, left, bottom, right = top_boxes[i]
                top     = max(0, np.floor(top).astype('int32'))
                left    = max(0, np.floor(left).astype('int32'))
                bottom  = min(image.size[1], np.floor(bottom).astype('int32'))
                right   = min(image.size[0], np.floor(right).astype('int32'))
                
                dir_save_path = "img_crop"
                if not os.path.exists(dir_save_path):
                    os.makedirs(dir_save_path)
                crop_image = image.crop([left, top, right, bottom])
                crop_image.save(os.path.join(dir_save_path, "crop_" + str(i) + ".png"), quality=95, subsampling=0)
                print("save crop_" + str(i) + ".png to " + dir_save_path)
        #---------------------------------------------------------#
        #   图像绘制
        #---------------------------------------------------------#
        for i, c in list(enumerate(top_label)):
            predicted_class = self.class_names[int(c)]
            box             = top_boxes[i]
            score           = top_conf[i]

            top, left, bottom, right = box

            top     = max(0, np.floor(top).astype('int32'))
            left    = max(0, np.floor(left).astype('int32'))
            bottom  = min(image.size[1], np.floor(bottom).astype('int32'))
            right   = min(image.size[0], np.floor(right).astype('int32'))

            # 名称和置信度
            label = '{} {:.2f}'.format(predicted_class, score)
            # 只显示名称
            # label = '{}'.format(predicted_class)
            draw = ImageDraw.Draw(image)
            label_size = draw.textsize(label, font)
            label = label.encode('utf-8')
            # print(label, top, left, bottom, right)
            
            if top - label_size[1] >= 0:
                text_origin = np.array([left, top - label_size[1]])
            else:
                text_origin = np.array([left, top + 1])

            for i in range(thickness):
                draw.rectangle([left + i, top + i, right - i, bottom - i], outline=self.colors[c])
            draw.rectangle([tuple(text_origin), tuple(text_origin + label_size)], fill=self.colors[c])
            draw.text(text_origin, str(label,'UTF-8'), fill=(0, 0, 0), font=font)
            del draw

        return image

    def get_FPS(self, image, test_interval):
        #---------------------------------------------------#
        #   计算输入图片的高和宽
        #---------------------------------------------------#
        image_shape = np.array(np.shape(image)[0:2])
        input_shape = get_new_img_size(image_shape[0], image_shape[1])
        #---------------------------------------------------------#
        #   在这里将图像转换成RGB图像，防止灰度图在预测时报错。
        #   代码仅仅支持RGB图像的预测，所有其它类型的图像都会转化成RGB
        #---------------------------------------------------------#
        image       = cvtColor(image)
        
        #---------------------------------------------------------#
        #   给原图像进行resize，resize到短边为600的大小上
        #---------------------------------------------------------#
        image_data  = resize_image(image, [input_shape[1], input_shape[0]])
        #---------------------------------------------------------#
        #   添加上batch_size维度
        #---------------------------------------------------------#
        image_data  = np.expand_dims(np.transpose(preprocess_input(np.array(image_data, dtype='float32')), (2, 0, 1)), 0)

        with torch.no_grad():
            images = torch.from_numpy(image_data)
            if self.cuda:
                images = images.to(self.device)

            #TODO
            # roi_cls_locs, roi_scores, rois, _ = self.model(images=images, mode='inference')
            output = self.model(images)
            #-------------------------------------------------------------#
            #   利用classifier的预测结果对建议框进行解码，获得预测框
            #-------------------------------------------------------------#
            # results = self.bbox_util.forward(roi_cls_locs, roi_scores, rois, image_shape, input_shape, 
            #                             nms_iou = self.nms_iou, confidence = self.confidence)
            results = self.bbox_util.forward(output, image_shape, input_shape, 
                                                    nms_iou = self.nms_iou, confidence = self.confidence)
            
        t1 = time.time()
        for _ in range(test_interval):
            with torch.no_grad():
                #TODO
                # roi_cls_locs, roi_scores, rois, _ = self.model(images=images, mode='inference')
                output = self.model(images)
                #-------------------------------------------------------------#
                #   利用classifier的预测结果对建议框进行解码，获得预测框
                #-------------------------------------------------------------#
                # results = self.bbox_util.forward(roi_cls_locs, roi_scores, rois, image_shape, input_shape, 
                #                             nms_iou = self.nms_iou, confidence = self.confidence)
                results = self.bbox_util.forward(output, image_shape, input_shape, 
                                                        nms_iou = self.nms_iou, confidence = self.confidence)
                
        t2 = time.time()
        tact_time = (t2 - t1) / test_interval
        return tact_time

    #---------------------------------------------------#
    #   检测图片
    #---------------------------------------------------#
    def get_map_txt(self, image_id, image, class_names, map_out_path):
        f = open(os.path.join(map_out_path, "detection-results/"+image_id+".txt"),"w")
        #---------------------------------------------------#
        #   计算输入图片的高和宽
        #---------------------------------------------------#
        image_shape = np.array(np.shape(image)[0:2])
        # input_shape = get_new_img_size(image_shape[0], image_shape[1])
        #---------------------------------------------------------#
        #   在这里将图像转换成RGB图像，防止灰度图在预测时报错。
        #   代码仅仅支持RGB图像的预测，所有其它类型的图像都会转化成RGB
        #---------------------------------------------------------#
        image       = cvtColor(image)
        
        #---------------------------------------------------------#
        #   给原图像进行resize，resize到短边为600的大小上
        #---------------------------------------------------------#
        image_data  = resize_image(image, [self.input_shape[1], self.input_shape[0]])
        #---------------------------------------------------------#
        #   添加上batch_size维度
        #---------------------------------------------------------#
        image_data  = np.expand_dims(np.transpose(preprocess_input(np.array(image_data, dtype='float32')), (2, 0, 1)), 0)

        with torch.no_grad():
            images = torch.from_numpy(image_data)
            if self.cuda:
                images = images.to(self.device)

            #TODO
            # roi_cls_locs, roi_scores, rois, _ = self.model(images=images, mode='inference')
            output = self.model(images)
            #-------------------------------------------------------------#
            #   利用classifier的预测结果对建议框进行解码，获得预测框
            #-------------------------------------------------------------#
            # results = self.bbox_util.forward(roi_cls_locs, roi_scores, rois, image_shape, input_shape, 
            #                             nms_iou = self.nms_iou, confidence = self.confidence)
            results = self.bbox_util.forward(output, image_shape, self.input_shape, 
                                                    nms_iou = self.nms_iou, confidence = self.confidence)
            #--------------------------------------#
            #   如果没有检测到物体，则返回原图
            #--------------------------------------#
            if len(results[0]) <= 0:
                return 

            top_label   = np.array(results[0][:, 5], dtype = 'int32')
            top_conf    = results[0][:, 4]
            top_boxes   = results[0][:, :4]
        
        for i, c in list(enumerate(top_label)):
            predicted_class = self.class_names[int(c)]
            box             = top_boxes[i]
            score           = str(top_conf[i])

            top, left, bottom, right = box
            if predicted_class not in self.class_names:
                continue

            f.write("%s %s %s %s %s %s\n" % (predicted_class, score[:6], str(int(left)), str(int(top)), str(int(right)),str(int(bottom))))

        f.close()
        return 
    
    # TODO 评估的计算量结果不准确！
    def get_model_complexity(self, input_shape):
        device  = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        model = self.model.to(device)
        summary(model,(3, input_shape[0], input_shape[1]))
        dummy_input     = torch.randn(1, 3, input_shape[0], input_shape[1]).to(device)
        flops, params   = profile(model.to(device), (dummy_input, ), verbose=False)
        flops           = flops * 2
        flops, params   = clever_format([flops, params], "%.3f")
        print('Total GFLOPS: %s' % (flops))
        print('Total params: %s' % (params))

    def detect_heatmap(self, image, heatmap_save_path):
        import cv2
        import matplotlib.pyplot as plt
        def sigmoid(x):
            y = 1.0 / (1.0 + np.exp(-x))
            return y
        #---------------------------------------------------------#
        #   在这里将图像转换成RGB图像，防止灰度图在预测时报错。
        #   代码仅仅支持RGB图像的预测，所有其它类型的图像都会转化成RGB
        #---------------------------------------------------------#
        image       = cvtColor(image)
        #---------------------------------------------------------#
        #   给图像增加灰条，实现不失真的resize
        #   也可以直接resize进行识别
        #---------------------------------------------------------#
        image_data  = resize_image(image, [self.input_shape[1], self.input_shape[0]])
        #---------------------------------------------------------#
        #   添加上batch_size维度
        #---------------------------------------------------------#
        image_data  = np.expand_dims(np.transpose(preprocess_input(np.array(image_data, dtype='float32')), (2, 0, 1)), 0)

        with torch.no_grad():
            images = torch.from_numpy(image_data)
            if self.cuda:
                images = images.to(self.device)
            #---------------------------------------------------------#
            #   将图像输入网络当中进行预测！
            #---------------------------------------------------------#
            # TODO 多尺度热力图预测，热力图颜色调整，可视化任意一层的热力图
            output = self.model(images, mode='backbone')
        heatmap = torch.sum(output, dim=1) 
        max_value = torch.max(heatmap)
        min_value = torch.min(heatmap)
        heatmap = (heatmap-min_value)/(max_value-min_value)*255
        heatmap = heatmap.cpu().numpy().astype(np.uint8).transpose(1,2,0)  # 提取热力图

        
        heatmap = cv2.resize(heatmap, self.input_shape,interpolation=cv2.INTER_LINEAR)  # 还原尺寸
        # 将矩阵转换为image类
        heatmap=cv2.applyColorMap(heatmap,cv2.COLORMAP_JET)
        heatimg = Image.fromarray(heatmap)
        
        # 将热力图叠加到原图上
        org_size = image.size
        heatimg = heatimg.resize(org_size)    #将热力图变回输入图像的尺寸
        plt.axis('off')
        plt.imshow(image)
        plt.imshow(heatimg, alpha=0.5)  # alpha为热力图的透明度

        # 保存叠加后的图形
        plt.savefig(heatmap_save_path, dpi=200, bbox_inches='tight', pad_inches = -0.1)
        plt.show()

    #TODO
    # def convert_to_onnx(self, simplify, model_path):
    #     import onnx
    #     self.generate(onnx=True)

    #     im                  = torch.zeros(1, 3, *self.input_shape).to('cpu')  # image size(1, 3, 512, 512) BCHW
    #     input_layer_names   = ["images"]
    #     output_layer_names  = ["output"]
        
    #     # Export the model
    #     print(f'Starting export with onnx {onnx.__version__}.')
    #     torch.onnx.export(self.model,
    #                     im,
    #                     f               = model_path,
    #                     verbose         = False,
    #                     opset_version   = 12,
    #                     training        = torch.onnx.TrainingMode.EVAL,
    #                     do_constant_folding = True,
    #                     input_names     = input_layer_names,
    #                     output_names    = output_layer_names,
    #                     dynamic_axes    = None)

    #     # Checks
    #     model_onnx = onnx.load(model_path)  # load onnx model
    #     onnx.checker.check_model(model_onnx)  # check onnx model

    #     # Simplify onnx
    #     if simplify:
    #         import onnxsim
    #         print(f'Simplifying with onnx-simplifier {onnxsim.__version__}.')
    #         model_onnx, check = onnxsim.simplify(
    #             model_onnx,
    #             dynamic_input_shape=False,
    #             input_shapes=None)
    #         assert check, 'assert check failed'
    #         onnx.save(model_onnx, model_path)

    #     print('Onnx model save as {}'.format(model_path))
